import os, json
import torch
import torchaudio
import psutil
import modelscope
from modelscope.pipelines import pipeline # type: ignore
from modelscope.utils.constant import Tasks # type: ignore
from typing import Literal
class FunASR:
    def __init__(
        self,
        funasr,
        model_dir,
        vad_model_dir = '',
        punc_model_dir = '',
        spk_model_dir = '',
        spkd_model_dir = '',
        cache_dir = '/dev/shm',
        deNoiseFunc = None,
        logLevel = "error"
    ):
        if vad_model_dir=='':
            vad_model_dir = None
        if punc_model_dir=='':
            punc_model_dir = None
        if spk_model_dir=='':
            spk_model_dir = None
        self.cache_dir = cache_dir
        self.logLevel = logLevel
        ngpu = 1
        device = "cuda" if torch.cuda.is_available() else "cpu"
        ncpu = psutil.cpu_count()
        ASRConfigs = dict(
            model=f"{model_dir}",
            vad_model_dir = vad_model_dir,
            punc_model_dir = punc_model_dir,
            spk_model_dir = spk_model_dir,
            ngpu=ngpu,
            ncpu=ncpu,
            device=device,
            weights_only=False,
            disable_update=True,
            disable_pbar=True if logLevel != 'debug' else False,
            disable_log=True if logLevel != 'debug' else False,
        )
        if vad_model_dir is None:
            del ASRConfigs['vad_model_dir']
        if punc_model_dir is None:
            del ASRConfigs['punc_model_dir']
        if spk_model_dir is None:
            del ASRConfigs['spk_model_dir']

        self.ASRModel = funasr.AutoModel(**ASRConfigs)
        self.SpeakerModel = None
        if spkd_model_dir != '' and spkd_model_dir is not None:
            self.SpeakerModel = pipeline(
                task='speaker-diarization',
                model=spkd_model_dir,
                device=device,
                disable_update=True,
                disable_pbar=True if logLevel != 'debug' else False,
                disable_log=True if logLevel != 'debug' else False,
            )
        self.deNoiseFunc = deNoiseFunc

    def speech(self, input, batch_size_s=300, output=None, target_rate = 16000, deNoise=True, useEmo=False, useSpeakerDiarization=False, output_timestamp=True,):
        result = {}
        if deNoise != True:
            deNoise = False
        if useEmo != True:
            useEmo = False
        if input is None or input == '':
            return {
                'code': 500,
                'msg': '请上传需要识别的音频文件',
            }
        if os.path.exists(input) == False:
            return {
                'code': 500,
                'msg': '需要识别的音频文件不存在',
            }
        if deNoise == True and self.deNoiseFunc is not None:
            self.deNoiseFunc(input, output, target_rate)
        res = [{"text":"","words":""}]
        try:
            res = self.ASRModel.generate(
                input=input,
                cache={},
                language="auto", # "zh", "en", "yue", "ja", "ko", "nospeech"
                use_itn=True,
                batch_size_s=batch_size_s,
                merge_vad=True,
                merge_length_s=15,
                output_timestamp=True,
                disable_update=True,
                ban_emo_unk=False
            )
        except:
            ...
        asrResult = dict(
            txt="".join(res[0]['words']),
            speakers=None
        )
        replaceString = {
            '"':'\"',
            '<|zh|>':'"lang":"zh",',
            '<|en|>':'"lang":"en",',
            '<|yue|>':'"lang":"yue",',
            '<|ja|>':'"lang":"ja",',
            '<|ko|>':'"lang":"ko",',
            '<|nospeech|>':'"lang":"nospeech",',
            '<|NEUTRAL|>':'"emo":"neutral",',
            '<|neutral|>':'"emo":"neutral",',
            '<|EMO_UNKNOWN|>':'"emo":"emo_unknown",',
            '<|emo_unknown|>':'"emo":"emo_unknown",',
            '<|HAPPY|>':'"emo":"happy",',
            '<|happy|>':'"emo":"happy",',
            '<|SAD|>':'"emo":"sad",',
            '<|sad|>':'"emo":"sad",',
            '<|ANGRY|>':'"emo":"angry",',
            '<|angry|>':'"emo":"angry",',
            '<|Speech|>':'',
            '<|speech|>':'',
            '<|woitn|>':'"txt":"',
            '<|withitn|>':'"txt":"',
        }
        if useEmo == True:
            txt=res[0]['text']
            for key, value in replaceString.items():
                txt=txt.replace(key,value)
            if txt.find('"txt":"') > -1:
                txt = '{'+ txt + '"}'
            asrResult=json.loads(txt)
        if useSpeakerDiarization == True and self.SpeakerModel is not None:
            result = self.SpeakerModel(input)
            asrResult['speakers'] = []
            for data in result['text']:
                sdata = []
                for d in data:
                    sdata.append(str(d))
                asrResult['speakers'].append(sdata)
        else:
            if 'speakers' in asrResult:
                del asrResult['speakers']
        asrResult = json.dumps(asrResult, ensure_ascii=False)
        return {
            "code": 200,
            "asr": asrResult
        }